"""
Triton Kernels for GPT-2 Optimization

这个包包含了用Triton实现的优化GPU kernels，用于替换PyTorch原生算子。
"""

from .attention import TritonGPT2Attention, triton_attention
from .layer_norm import TritonLayerNorm
from .mlp import TritonGPT2MLP, TritonGELU, TritonFusedMLP
from .utils import TritonModuleBase

__all__ = [
    'TritonGPT2Attention',
    'TritonLayerNorm',
    'TritonGPT2MLP',
    'TritonGELU',
    'TritonFusedMLP',
    'TritonModuleBase',
    'triton_attention',
]

